from kuai.sim.energy import Energy


class ExternalForces(Energy):
    FORCE_X = 1
    FORCE_Y = 2
    FORCE_Z = 4
    FORCE_XY = FORCE_X | FORCE_Y
    FORCE_YZ = FORCE_Y | FORCE_Z
    FORCE_XZ = FORCE_X | FORCE_Z
    FORCE_XYZ = FORCE_X | FORCE_Y | FORCE_Z  

    def __init__(self, name, direction=FORCE_XYZ):
        Energy.__init__(self, name)
        self.direction = direction

    def energies(self, mol, ff, energies):
        pass
    
    def forces(self, mol, ff, energies, forces):
        raise NotImplementedError()
    
    def hessian(self, mol, ff, energies, forces, hessian):
        self.forces(mol, ff, energies, forces)

    def rank(self):
        raise NotImplementedError()


class RandomUniformForces(ExternalForces):
    def __init__(self, direction):
        ExternalForces.__init__(self, 'RandomUniformForces', direction)
        
    def forces(self, mol, ff, energies, forces):
        from kuai.kuaiext import random_uniform
        natoms = len(mol.atoms)
        if self.direction & ExternalForces.FORCE_X:
            x = forces[0:len(forces):3]
            x += random_uniform()
        if self.direction & ExternalForces.FORCE_Y:
            y = forces[1:len(forces):3]
            y += forces[1:len(forces):3]
        if self.direction & ExternalForces.FORCE_Z:
            z = forces[2:len(forces):3]
            z += forces[2:len(forces):3]
        
